﻿/*
 * @author: wizardc
 */

using System.Collections.Generic;
using System.Linq;
using Dou.GPU;
using Dou.Utils;
using UnityEditor;
using UnityEngine;

namespace DouEditor.GPU
{
    /// <summary>
    /// GPU 动画数据录制
    /// </summary>
    public static class GPUAnimRecord
    {
        /// <summary>
        /// 录制
        /// </summary>
        public static bool Record(GameObject asset, out GPUAnimData animData, out Texture2D animTex, out Mesh animMesh)
        {
            var assetPath = AssetDatabase.GetAssetPath(asset);
            var go = GameObject.Instantiate(asset);
            
            if (!GetClips(assetPath, go, out var skinnedMesh, out var clips))
            {
                animData = null;
                animTex = null;
                animMesh = null;
                return false;
            }

            go.transform.SetParent(null);
            go.transform.position = Vector3.zero;
            go.transform.rotation = Quaternion.identity;
            go.transform.localScale = Vector3.one;

            animData = new GPUAnimData();
            animTex = CreateTexture(go, animData, skinnedMesh, clips);
            animMesh = CreateAnimMesh(skinnedMesh);
            GameObject.DestroyImmediate(go);
            
            return true;
        }

        private static bool GetClips(string assetPath, GameObject go, out SkinnedMeshRenderer skinnedMesh, out AnimationClip[] clips)
        {
            var meshes = go.GetComponentsInChildren<SkinnedMeshRenderer>();
            if (meshes.Length == 0)
            {
                skinnedMesh = null;
                clips = null;
                Log.Warning($"\"{assetPath}\"没有找到蒙皮网格，GPU动画录制失败！");
                return false;
            }
            if (meshes.Length > 1)
            {
                Log.Warning($"\"{assetPath}\"存在多个蒙皮网格，只录制第一个！");
            }
            skinnedMesh = meshes[0];
            
            var animator = go.GetComponent<Animator>();
            if (animator != null)
            {
                clips = animator.runtimeAnimatorController.animationClips;
            }
            else
            {
                var animation = go.GetComponent<Animation>();
                if (animation == null)
                {
                    skinnedMesh = null;
                    clips = null;
                    Log.Warning($"\"{assetPath}\"没有找到动画组件，GPU动画录制失败！");
                    return false;
                }
                var count = animation.GetClipCount();
                clips = new AnimationClip[count];
                var index = 0;
                foreach (AnimationState clip in animation) {
                    clips[index++] = animation.GetClip(clip.name);
                }
            }
            return true;
        }

        private static Texture2D CreateTexture(GameObject go, GPUAnimData animData, SkinnedMeshRenderer skinnedMesh, AnimationClip[] clips)
        {
            animData.boneCount = skinnedMesh.bones.Length;
            var clipList = new List<GPUAnimClipData>();
            var offset = 0;
            var matrixList = new List<List<List<Matrix4x4>>>();
            foreach (var clip in clips) {
                RecordClip(go, clipList, skinnedMesh, clip, ref offset, matrixList);
            }
            animData.frameCount = matrixList.Sum(clip => clip.Count);
            animData.clips = clipList.ToArray();

            var texture = new Texture2D(animData.frameCount, animData.boneCount * 3, TextureFormat.RGBAHalf, false, true);
            texture.filterMode = FilterMode.Point;
            var col = 0;
            foreach (var frames in matrixList) {
                foreach (var bones in frames) {
                    for (var i = 0; i < bones.Count; i++) {
                        var matrix = bones[i];
                        var line = i * 3;
                        texture.SetPixel(col, line + 0, matrix.GetRow(0));
                        texture.SetPixel(col, line + 1, matrix.GetRow(1));
                        texture.SetPixel(col, line + 2, matrix.GetRow(2));
                    }
                    col++;
                }
            }
            texture.Apply();
            return texture;
        }

        private static void RecordClip(GameObject go, List<GPUAnimClipData> animData, SkinnedMeshRenderer skinnedMesh, AnimationClip clip, ref int offset, List<List<List<Matrix4x4>>> matrixList)
        {
            var fps = clip.frameRate;
            var len = clip.length;
            var frames = Mathf.CeilToInt(fps * len);
            var frameTime = 1f / fps;
            
            animData.Add(new GPUAnimClipData
            {
                name = clip.name,
                startIndex = offset,
                frameCount = frames,
                frameTime = frameTime,
                isLoop = clip.isLooping,
            });
            
            offset += frames;

            var time = 0f;
            var frameMatrixList = new List<List<Matrix4x4>>();
            for (int i = 0; i < frames; i++)
            {
                RecordFrame(time, go, skinnedMesh, clip, frameMatrixList);
                time += frameTime;
            }
            
            matrixList.Add(frameMatrixList);
        }

        private static void RecordFrame(float time, GameObject go, SkinnedMeshRenderer skinnedMesh, AnimationClip clip, List<List<Matrix4x4>> frameMatrixList)
        {
            clip.SampleAnimation(go, time);
            var bones = skinnedMesh.bones;
            var poses = skinnedMesh.sharedMesh.bindposes;
            var boneMatrix = new List<Matrix4x4>();
            for (var i = 0; i < bones.Length; i++) {
                var matrix = bones[i].localToWorldMatrix * poses[i];
                boneMatrix.Add(matrix);
            }
            frameMatrixList.Add(boneMatrix);
        }

        private static Mesh CreateAnimMesh(SkinnedMeshRenderer skinnedMesh)
        {
            var origin = skinnedMesh.sharedMesh;
            var mesh = new Mesh {
                name = origin.name,
                vertices = origin.vertices,
                uv = origin.uv,
                normals = origin.normals,
                tangents = origin.tangents,
                colors = origin.colors,
                triangles = origin.triangles,
                bounds = origin.bounds
            };
            mesh.RecalculateBounds();
            mesh.RecalculateNormals();
            mesh.RecalculateTangents();
            
            var boneWeights = origin.boneWeights;
            var boneLen = boneWeights.Length;
            var idxList = new List<Vector4>();
            var weightList = new List<Vector4>();
            for (var i = 0; i < boneLen; i++) {
                var weight = boneWeights[i];
                var tempBone = new Vector4 {
                    x = weight.boneIndex0, y = weight.boneIndex1, z = weight.boneIndex2, w = weight.boneIndex3
                };
                idxList.Add(tempBone);
                var tempWeight = new Vector4 {
                    x = weight.weight0, y = weight.weight1, z = weight.weight2, w = weight.weight3
                };
                weightList.Add(tempWeight);
            }
            mesh.SetUVs(1, idxList);
            mesh.SetUVs(2, weightList);
            return mesh;
        }
    }
}
